refactor(pt): reuse dpmodel NeighborStatOP #5137
Conversation
This class does not need torchscript, so we can reuse dpmodel codes to reduce redundancy.
for more information, see https://pre-commit.ci
📝 WalkthroughWalkthroughDevice-consistency fixes and a small refactor: array creations in neighbor-list and neighbor-stat utilities now explicitly allocate on the input device; PyTorch-side NeighborStat now delegates to the dpmodel NeighborStatOP implementation. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
📜 Recent review detailsConfiguration used: Repository UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
🧰 Additional context used🧠 Learnings (1)📓 Common learnings⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (37)
🔇 Additional comments (3)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
deepmd/pt/utils/neighbor_stat.py (1)
24-25: Minor: Update docstring to reflect implementation change.The docstring states "Neighbor statistics using pure NumPy" but the implementation now uses
NeighborStatOPfromdpmodel, which operates viaarray_api_compatand can work with various backends including PyTorch tensors.📝 Suggested docstring update
class NeighborStat(BaseNeighborStat): - """Neighbor statistics using pure NumPy. + """Neighbor statistics using PyTorch. Parameters
📜 Review details
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
deepmd/dpmodel/utils/neighbor_stat.pydeepmd/dpmodel/utils/nlist.pydeepmd/pt/utils/neighbor_stat.py
🧰 Additional context used
🧠 Learnings (2)
📓 Common learnings
Learnt from: njzjz
Repo: deepmodeling/deepmd-kit PR: 4144
File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246
Timestamp: 2024-09-19T04:25:12.408Z
Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905.
Learnt from: njzjz
Repo: deepmodeling/deepmd-kit PR: 4144
File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246
Timestamp: 2024-10-08T15:32:11.479Z
Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905.
📚 Learning: 2024-10-08T15:32:11.479Z
Learnt from: njzjz
Repo: deepmodeling/deepmd-kit PR: 4160
File: deepmd/dpmodel/utils/env_mat.py:52-64
Timestamp: 2024-10-08T15:32:11.479Z
Learning: Negative indices in `nlist` are properly handled by masking later in the computation, so they do not cause issues in indexing operations.
Applied to files:
deepmd/dpmodel/utils/nlist.py
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (42)
- GitHub Check: CodeQL analysis (python)
- GitHub Check: Agent
- GitHub Check: Test Python (10, 3.10)
- GitHub Check: Test Python (11, 3.10)
- GitHub Check: Test Python (12, 3.13)
- GitHub Check: Test Python (9, 3.10)
- GitHub Check: Test Python (12, 3.10)
- GitHub Check: Test Python (5, 3.13)
- GitHub Check: Test Python (11, 3.13)
- GitHub Check: Test Python (6, 3.10)
- GitHub Check: Test Python (6, 3.13)
- GitHub Check: Test Python (10, 3.13)
- GitHub Check: Test Python (8, 3.13)
- GitHub Check: Test Python (3, 3.13)
- GitHub Check: Test Python (9, 3.13)
- GitHub Check: Test Python (8, 3.10)
- GitHub Check: Test Python (4, 3.10)
- GitHub Check: Test Python (1, 3.10)
- GitHub Check: Analyze (python)
- GitHub Check: Test Python (1, 3.13)
- GitHub Check: Test Python (2, 3.10)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Test Python (2, 3.13)
- GitHub Check: Test Python (3, 3.10)
- GitHub Check: Test Python (5, 3.10)
- GitHub Check: Test Python (7, 3.10)
- GitHub Check: Test Python (7, 3.13)
- GitHub Check: Test Python (4, 3.13)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Test C++ (true, true, true, false)
- GitHub Check: Test C++ (false, false, false, true)
- GitHub Check: Test C++ (true, false, false, true)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Test C++ (false, true, true, false)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
🔇 Additional comments (8)
deepmd/dpmodel/utils/neighbor_stat.py (1)
90-90: LGTM! Device-aware array creation for cross-backend compatibility.The explicit
device=array_api_compat.device(diff)ensures the identity matrix is allocated on the same device as the input-deriveddifftensor. This is essential for PyTorch backend usage where device mismatches would cause runtime errors.deepmd/dpmodel/utils/nlist.py (5)
120-122: LGTM! Device-consistent identity matrix for zero-distance correction.The
device=array_api_compat.device(diff)ensures the eye matrix is on the same device as the computed differences, preventing device mismatch errors on GPU backends.
132-154: LGTM! Device-aware padding for neighbor list extension.Both padding arrays correctly inherit the device from their source arrays (
rrandnlistrespectively), ensuring consistent device placement when extending the neighbor list.
238-243: LGTM! Device-aware padding in multiple neighbor list builder.The padding array correctly uses
device=array_api_compat.device(nlist)to ensure it resides on the same device as the input neighbor list.
300-305: LGTM! Device-aware atom index creation.The
aidxarray correctly usesdevice=array_api_compat.device(atype)to ensure the index array resides on the same device as the atom types.
317-351: LGTM! Consistent device propagation for ghost coordinate generation.All intermediate arrays (
xi,yi,zi, and the unit vectors for outer products) are correctly allocated on the same device as the inputcoordtensor. The unit vectors appropriately use the device from their corresponding index arrays, maintaining device consistency throughout the ghost extension computation.deepmd/pt/utils/neighbor_stat.py (2)
92-98: LGTM! Correct tensor conversion flow.The conversion chain from NumPy arrays to PyTorch tensors on
DEVICE, through theNeighborStatOPcall, and back to NumPy via.detach().cpu().numpy()is correct.
19-21: Verify array-api-compat torch tensor compatibility through existing test suite.The refactoring to reuse
NeighborStatOPfromdpmodelis sound—the implementation correctly usesarray_api_compat.array_namespace()to abstract array operations (line 72 of dpmodel/utils/neighbor_stat.py), and the PyTorch wrapper properly converts tensors to/from numpy. However, ensure that the pytorch backend tests (test_neighbor_stat_ptin source/tests/consistent/test_neighbor_stat.py) pass to confirm thatarray-api-compathandles torch tensors correctly across all operations, particularlyxp.booldtype specification and device handling withxp.eye().
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #5137 +/- ##
==========================================
- Coverage 82.14% 82.13% -0.01%
==========================================
Files 709 709
Lines 72511 72480 -31
Branches 3615 3615
==========================================
- Hits 59565 59534 -31
- Misses 11782 11783 +1
+ Partials 1164 1163 -1 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
github-merge-queue
bot
removed this pull request from the merge queue due to failed status checks
3 participants
This class does not need torchscript, so we can reuse dpmodel codes to reduce redundancy.
Summary by CodeRabbit
Bug Fixes
Refactor
✏️ Tip: You can customize this high-level summary in your review settings.